"""
# labels.csv已知
# 病灶最大最小半径，最大最小Z位置统计
# 扩充数据集1：假设当前病灶为球形，以当前切片为已知信息，向上扩展5张，向下扩展5张
# 扩充数据集2：经过观察，对病灶半径小于10的，上下扩充一张；对病灶半径大于10的，上下扩充三张
# 以扩充数据集为依据，提取CT和PET dicom图像转存，加速之后预处理
# 统一csv文件中的身高格式
# 生成三维立方体模型输入

"""


import os
import csv
import shutil
import numpy as np
import pandas as pd
import pydicom as pydic
from threading import Thread

# 统计最大、最小病灶半径R
# 统计最大、最小病灶位置Z
def statisticRZ():
    myfile = open('D:/lung_cancer/data/data_augmentation/labels.csv')
    labels = []
    lines = csv.reader(myfile)
    for line in lines:
        labels.append(line)

    # 设置初始值
    minR = 300
    maxR = 0
    minZ = 1000
    maxZ = 0

    for line in labels[1:]:
        R = int(line[5])
        Z = int(line[2])

        if R > maxR:
            maxR = R
        if R<minR:
            minR = R
        if Z > maxZ:
            maxZ = Z
        if Z < minZ:
            minZ = Z

    print('minR: %d, maxR: %d, minZ: %d, maxZ: %d'%(minR, maxR, minZ, maxZ))



# 扩充labels.csv, 生成aug_labels.csv
# 经统计得到最小病灶半径为5，我们上下各取五张，病灶半径变换方式查看visio图
def augmentData1():
    myfile = open('D:/lung_cancer/data/data_augmentation/labels.csv')
    labels = []
    lines = csv.reader(myfile)
    for line in lines:
        labels.append(line)

    infos = []

    for line in labels[1:]:
        z = int(line[2])
        r = int(line[5])
        infos.append(line)
        for i in range(5):  # 上下各扩展五张
            up_line = line[:]  # 向上扩展
            down_line = line[:]  # 向下扩展
            up_line[2] = z-i-1
            up_line[5] = r-i
            down_line[2] = z+i+1
            down_line[5] = r-i
            infos.append(up_line)
            infos.append(down_line)

    # 保存扩展标签
    df = pd.DataFrame(infos, columns=labels[0])
    df.to_csv('D:/lung_cancer/data/data_augmentation/augement_labels.csv', index=False)



# 扩充labels.csv, 生成aug2_labels.csv

def augmentData2():
    myfile = open('D:/lung_cancer/data/data_augmentation/labels.csv')
    labels = []
    lines = csv.reader(myfile)
    for line in lines:
        labels.append(line)

    infos = []

    for line in labels[1:]:
        z = int(line[2])
        r = int(line[5])
        infos.append(line)

        if z < 10:
            up_line = line[:]
            down_line = line[:]
            up_line[2] = z-1
            down_line[2] = z+1

            infos.append(up_line)
            infos.append(down_line)

        else:
            for i in range(3):  # 上下各扩展三张
                up_line = line[:]  # 向上扩展
                down_line = line[:]  # 向下扩展
                up_line[2] = z-i-1
                up_line[5] = r-i
                down_line[2] = z+i+1
                down_line[5] = r-i
                infos.append(up_line)
                infos.append(down_line)

    # 保存扩展标签
    df = pd.DataFrame(infos, columns=labels[0])
    df.to_csv('D:/lung_cancer/data/data_augmentation/augement2_labels.csv', index=False)



# 得到指定Instance Number的dicom图像名字
def getName(origin_path, SeriesDescription, InstanceNumber):
    names = os.listdir(origin_path)

    for name in names:
        slice = pydic.read_file(origin_path + '/' + name)
        if slice.SeriesDescription == SeriesDescription and slice.InstanceNumber == InstanceNumber:
            return name

    print('------------------------------------------', origin_path.split('/')[-1], '-----------------------------')

    return 'no'



# 分别读取aug_labels.csv和aug2_labels.csv的CT和PET信息，生成 all_data1.csv和all_2_data1.csv
def save_info():

    num = 1

    columns = ['patientID', 'z', 'x', 'y', 'r', 'cancer_type', 'ct_count', 'pet_count', 'ct_size', 'pet_size',
               'ct_x_spacing', 'ct_y_spacing', 'ct_z_spacing', 'pet_x_spacing', 'pet_y_spacing', 'pet_z_spacing',
               'ct_intercept', 'pet_intercept', 'ct_slope', 'pet_slope', 'patientWeight', 'patientSex', 'patientAge',
               'patientSize', 'TotalDose', 't0', 't1', 'HalfLife', 'DecayFactor', 'CT_SeriesDescription',
               'PET_SeriesDescription', 'origin_dir', 'CT_origin_path', 'PET_origin_path']

    # myfile = open('D:/lung_cancer/data/data_augmentation/augement_labels.csv')
    myfile = open('D:/lung_cancer/data/data_augmentation/augement2_labels.csv')

    labels = []
    lines = csv.reader(myfile)
    for line in lines:
        labels.append(line)

    infos = []
    for line in labels[1:]:

        path = 'H:/'+line[9]
        names = os.listdir(path)
        ct_slices = []
        pet_slices = []

        for name in names:
            slice = pydic.read_file(path + '/' + name)
            if slice.SeriesDescription == str(line[7]):
                ct_slices.append(slice)
            elif slice.SeriesDescription == str(line[8]):
                pet_slices.append(slice)
        ct_slices.sort(key=lambda x: int(x.InstanceNumber))  # 默认升序排列
        pet_slices.sort(key=lambda y: int(y.InstanceNumber))  # 默认升序排列

        ct_count = len(ct_slices)
        pet_count = len(pet_slices)


        ct_z_spacing = np.abs(ct_slices[0].SliceLocation - ct_slices[1].SliceLocation)
        pet_z_spacing = np.abs(pet_slices[0].SliceLocation - pet_slices[1].SliceLocation)

        patientID = int(line[1])
        z = int(line[2])
        x = int(line[3])
        y = int(line[4])
        r = int(line[5])
        cancer_type = int(line[6])

        CT_SeriesDescription = line[7]
        PET_SeriesDescription = line[8]

        # 得到原始地址
        CT_origin_path = getName(path, CT_SeriesDescription, z)
        PET_origin_path = getName(path, PET_SeriesDescription, z)

        # 方法一：用附近的代替
        # while CT_origin_path == 'no' or PET_origin_path == 'no':
        #     z = z+1
        #     CT_origin_path = getName(path, CT_SeriesDescription, z)
        #     PET_origin_path = getName(path, PET_SeriesDescription, z)

        # 方法二：直接跳过
        if CT_origin_path == 'no' or PET_origin_path == 'no':
            continue

        CT_origin_path = line[9]+'/'+getName(path, CT_SeriesDescription, z)
        PET_origin_path = line[9]+'/'+getName(path, PET_SeriesDescription, z)


        # 可能因为ct和pet的Instance Number不一定从0开始，所以不能直接用上文的索引得到ct和pet
        ct = pydic.read_file('H:/'+CT_origin_path)
        pet = pydic.read_file('H:/'+PET_origin_path)

        ct_size = ct.Rows
        pet_size = pet.Rows
        ct_x_spacing = ct.PixelSpacing[0]
        ct_y_spacing = ct.PixelSpacing[1]
        pet_x_spacing = pet.PixelSpacing[0]
        pet_y_spacing = pet.PixelSpacing[1]
        ct_intercept = ct.RescaleIntercept
        pet_intercept = pet.RescaleIntercept
        ct_slope = ct.RescaleSlope
        pet_slope = pet.RescaleSlope

        patientAge = int(ct[0x10, 0x1010].value[:-1])
        patientWeight = float(pet[0x10, 0x1030].value)
        patientSize = float(pet[0x10, 0x1020].value) * 100
        Sex = pet[0x10, 0x40].value
        patientSex = 0  # 男性为0， 女性为1
        if (Sex == 'F'):
            patientSex = 1
        DecayFactor = pet.DecayFactor

        a = pet.RadiopharmaceuticalInformationSequence
        TotalDose = a[0].RadionuclideTotalDose
        HalfLife = a[0].RadionuclideHalfLife
        t0 = a[0].RadiopharmaceuticalStartTime
        t1 = pet.AcquisitionTime


        origin_dir = line[9]



        infos.append([patientID, z, x, y, r, cancer_type, ct_count, pet_count, ct_size, pet_size, ct_x_spacing,
                      ct_y_spacing, ct_z_spacing, pet_x_spacing, pet_y_spacing, pet_z_spacing, ct_intercept,
                      pet_intercept, ct_slope, pet_slope, patientWeight, patientSex, patientAge,patientSize,
                      TotalDose, t0, t1, HalfLife, DecayFactor, CT_SeriesDescription, PET_SeriesDescription, origin_dir,
                      CT_origin_path, PET_origin_path])
        print(num, '--->', patientID, 'is ok')
        num = num+1
    df = pd.DataFrame(infos, columns=columns)
    # df.to_csv('D:/lung_cancer/data/data_augmentation/all_data1.csv', index=False)
    df.to_csv('D:/lung_cancer/data/data_augmentation/all_2_data1.csv', index=False)

# 生成的身高有些问题，重新格式化身高, 生成新的all_data1.csv和all_2_data1.csv
def standardPatientSize():
    # myfile = open('D:/lung_cancer/data/data_augmentation/all_data1.csv')
    myfile = open('D:/lung_cancer/data/data_augmentation/all_2_data1.csv')

    labels = []
    lines = csv.reader(myfile)
    for line in lines:
        labels.append(line)

    infos = []
    for line in labels[1:]:
        patientSize = int(float(line[23]))
        print(patientSize)
        if patientSize > 999:
            print(line[0])
            line[23] = int(patientSize/100)
        infos.append(line)

    df = pd.DataFrame(infos, columns=labels[0])
    # df.to_csv('D:/lung_cancer/data/data_augmentation/all_data1.csv', index=False)
    df.to_csv('D:/lung_cancer/data/data_augmentation/all_2_data1.csv', index=False)


# 分别利用all_data1.csv和all_2_data1.csv中的文件地址，把要用到的CT和PET单独保存
def extract_origin_CT_PET():
    # data = pd.read_csv('D:/lung_cancer/data/data_augmentation/all_data1.csv')
    data = pd.read_csv('D:/lung_cancer/data/data_augmentation/all_2_data1.csv')
    CT_path = data['CT_origin_path']
    PET_path = data['PET_origin_path']
    patientID = data['patientID']
    print(patientID[len(patientID)-1])

    for i in range(len(patientID)):
        print(i, '----->', patientID[i])
        CT_save_path = 'D:/lung_cancer/data/data_augmentation/origin_2_data/'+str(patientID[i])+'/CT/'
        CT_origin_path = 'H:/'+CT_path[i]

        PET_save_path = 'D:/lung_cancer/data/data_augmentation/origin_2_data/'+str(patientID[i])+'/PET/'
        PET_origin_path = 'H:/' + PET_path[i]

        if (not os.path.exists(CT_save_path)):
            os.makedirs(CT_save_path)
        shutil.copy(CT_origin_path, CT_save_path)

        if (not os.path.exists(PET_save_path)):
            os.makedirs(PET_save_path)
        shutil.copy(PET_origin_path, PET_save_path)


# 读取二维切片组合为三维输入
def generate_3D_img(path):
    # path = 'D:/lung_cancer/data/data_augmentation/Slice/3107/CTSlice/'
    ctList = os.listdir(path)
    ctList.sort()

    imglist = []
    for i in range(len(ctList)):
        img = np.load(path+ctList[i])
        imglist.append(img)

    # 对二维图片多于8张的去掉，少的补0
    num = len(imglist)
    if num > 8:
        imglist = imglist[:8]
    elif num < 8:
        for i in range(8-num):
            imglist.append(np.zeros((512, 512)))

    imgs = np.asarray(imglist)

    # z轴变到最后
    imgs = imgs.transpose((1, 2, 0))

    return imgs



if __name__ == '__main__':
    # statisticRZ()

    # augmentData1()
    # augmentData2()

    # save_info()
    # standardPatientSize()

    extract_origin_CT_PET()

    # generate_3D_img()